import { useCall, usePubSub, useWillMountEffect } from '@/components';
import { defaultShouldUpdate, mergeState } from '@/utils';
import { get } from '@/utils/fp';
import React, { createContext, forwardRef, useContext, useImperativeHandle, useMemo, useState } from 'react';

const StateContext = createContext({});

StateProvider = forwardRef(StateProvider);

export function StateProvider({ children, initialValue, shouldUpdate }, ref) {
  const pubsub = usePubSub(shouldUpdate, initialValue);

  const setState = useCall((value) => {
    const newValue = mergeState(pubsub.value, typeof value === 'function' ? value(pubsub.value) : value);
    pubsub.update(newValue);
  });

  const getState = useCall((...keys) => get(keys)(pubsub.value));

  useImperativeHandle(
    ref,
    () => ({
      pubsub,
      getState,
      setState,
    }),
    [],
  );

  const context = useMemo(() => ({ pubsub, setState }), []);

  return <StateContext.Provider value={context} children={children} />;
}

export function useSelector(selector = [], shouldUpdate = defaultShouldUpdate) {
  const { pubsub } = useContext(StateContext);
  const [state, setState] = useState(get(selector)(pubsub.value ?? null));
  useWillMountEffect(pubsub.sub, (value) => {
    const newState = get(selector)(value ?? null);
    setState((prevState) => {
      if (shouldUpdate(prevState, newState)) return newState;
      return prevState;
    });
  });

  return state;
}

export function useUpdater() {
  const { setState } = useContext(StateContext);
  return setState;
}
